// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include <zencore/filesystem.h>
#include <zencore/zencore.h>
#include <zenutil/basicfile.h>

#include <unordered_map>
#include <unordered_set>

namespace zen {

//////////////////////////////////////////////////////////////////////////

struct BlockStoreLocation
{
	uint32_t BlockIndex;
	uint32_t Offset;
	uint32_t Size;

	inline auto operator<=>(const BlockStoreLocation& Rhs) const = default;
};

#pragma pack(push)
#pragma pack(1)

struct BlockStoreDiskLocation
{
	constexpr static uint32_t MaxBlockIndexBits = 20;
	constexpr static uint32_t MaxOffsetBits		= 28;
	constexpr static uint32_t MaxBlockIndex		= (1ul << BlockStoreDiskLocation::MaxBlockIndexBits) - 1ul;
	constexpr static uint32_t MaxOffset			= (1ul << BlockStoreDiskLocation::MaxOffsetBits) - 1ul;

	BlockStoreDiskLocation(const BlockStoreLocation& Location, uint32_t OffsetAlignment)
	{
		Init(Location.BlockIndex, Location.Offset / OffsetAlignment, Location.Size);
	}

	BlockStoreDiskLocation() = default;

	inline BlockStoreLocation Get(uint32_t OffsetAlignment) const
	{
		uint64_t PackedOffset = 0;
		memcpy(&PackedOffset, &m_Offset, sizeof m_Offset);
		return {.BlockIndex = static_cast<uint32_t>(PackedOffset >> MaxOffsetBits),
				.Offset		= static_cast<uint32_t>((PackedOffset & MaxOffset) * OffsetAlignment),
				.Size		= GetSize()};
	}

	inline uint32_t GetBlockIndex() const
	{
		uint64_t PackedOffset = 0;
		memcpy(&PackedOffset, &m_Offset, sizeof m_Offset);
		return static_cast<std::uint32_t>(PackedOffset >> MaxOffsetBits);
	}

	inline uint32_t GetOffset(uint32_t OffsetAlignment) const
	{
		uint64_t PackedOffset = 0;
		memcpy(&PackedOffset, &m_Offset, sizeof m_Offset);
		return static_cast<uint32_t>((PackedOffset & MaxOffset) * OffsetAlignment);
	}

	inline uint32_t GetSize() const { return m_Size; }

	inline auto operator<=>(const BlockStoreDiskLocation& Rhs) const = default;

private:
	inline void Init(uint32_t BlockIndex, uint64_t Offset, uint64_t Size)
	{
		ZEN_ASSERT(BlockIndex <= MaxBlockIndex);
		ZEN_ASSERT(Offset <= MaxOffset);
		ZEN_ASSERT(Size <= std::numeric_limits<std::uint32_t>::max());

		m_Size				  = static_cast<uint32_t>(Size);
		uint64_t PackedOffset = (static_cast<uint64_t>(BlockIndex) << MaxOffsetBits) + Offset;
		memcpy(&m_Offset[0], &PackedOffset, sizeof m_Offset);
	}

	uint32_t m_Size;
	uint8_t	 m_Offset[6];
};

static_assert(sizeof(BlockStoreDiskLocation) == 10);

#pragma pack(pop)

struct BlockStoreFile : public RefCounted
{
	explicit BlockStoreFile(const std::filesystem::path& BlockPath);
	~BlockStoreFile();
	const std::filesystem::path& GetPath() const;
	void						 Open();
	void						 Create(uint64_t InitialSize);
	void						 MarkAsDeleteOnClose();
	uint64_t					 FileSize();
	IoBuffer					 GetChunk(uint64_t Offset, uint64_t Size);
	void						 Read(void* Data, uint64_t Size, uint64_t FileOffset);
	void						 Write(const void* Data, uint64_t Size, uint64_t FileOffset);
	void						 Flush();
	BasicFile&					 GetBasicFile();
	void StreamByteRange(uint64_t FileOffset, uint64_t Size, std::function<void(const void* Data, uint64_t Size)>&& ChunkFun);
	bool IsOpen() const;

private:
	const std::filesystem::path m_Path;
	IoBuffer					m_IoBuffer;
	BasicFile					m_File;
	uint64_t					m_CachedFileSize = 0;
};

class BlockStoreCompactState;

class BlockStore
{
public:
	BlockStore();
	~BlockStore();

	struct ReclaimSnapshotState
	{
		std::unordered_set<uint32_t> m_ActiveWriteBlocks;
		std::unordered_set<uint32_t> m_BlockIndexes;
	};

	typedef std::vector<std::pair<size_t, BlockStoreLocation>> MovedChunksArray;
	typedef std::vector<size_t>								   ChunkIndexArray;

	typedef std::function<void(const MovedChunksArray& MovedChunks, const ChunkIndexArray& RemovedChunks)> ReclaimCallback;
	typedef std::function<bool(const MovedChunksArray& MovedChunks, uint64_t FreedDiskSpace)>			   CompactCallback;
	typedef std::function<uint64_t()>																	   ClaimDiskReserveCallback;
	typedef std::function<bool(size_t ChunkIndex, const void* Data, uint64_t Size)>						   IterateChunksSmallSizeCallback;
	typedef std::function<bool(size_t ChunkIndex, BlockStoreFile& File, uint64_t Offset, uint64_t Size)>   IterateChunksLargeSizeCallback;
	typedef std::function<void(const BlockStoreLocation& Location)>										   WriteChunkCallback;
	typedef std::function<bool(uint32_t BlockIndex, std::span<const size_t> ChunkIndexes)>				   IterateChunksCallback;

	struct BlockUsageInfo
	{
		uint64_t DiskUsage;
		uint32_t EntryCount;
	};
	typedef std::unordered_map<uint32_t, BlockUsageInfo> BlockUsageMap;
	typedef std::unordered_map<uint32_t, uint32_t>		 BlockEntryCountMap;

	void Initialize(const std::filesystem::path& BlocksBasePath, uint64_t MaxBlockSize, uint64_t MaxBlockCount);

	struct BlockIndexSet
	{
		void					  Add(uint32_t BlockIndex);
		std::span<const uint32_t> GetBlockIndices() const { return BlockIndexes; }

	private:
		std::vector<uint32_t> BlockIndexes;
	};

	// Ask the store to create empty blocks for all locations that does not have a block
	// Remove any block that is not referenced
	void			   SyncExistingBlocksOnDisk(const BlockIndexSet& KnownLocations);
	BlockEntryCountMap GetBlocksToCompact(const BlockUsageMap& BlockUsage, uint32_t BlockUsageThresholdPercent);

	void Close();

	void WriteChunk(const void* Data, uint64_t Size, uint32_t Alignment, const WriteChunkCallback& Callback);

	typedef std::function<void(std::span<BlockStoreLocation> Locations)> WriteChunksCallback;
	void WriteChunks(std::span<IoBuffer> Datas, uint32_t Alignment, const WriteChunksCallback& Callback);

	IoBuffer TryGetChunk(const BlockStoreLocation& Location) const;
	void	 Flush(bool ForceNewBlock);

	ReclaimSnapshotState GetReclaimSnapshotState();
	void				 ReclaimSpace(
						const ReclaimSnapshotState&			   Snapshot,
						const std::vector<BlockStoreLocation>& ChunkLocations,
						const ChunkIndexArray&				   KeepChunkIndexes,
						uint32_t							   PayloadAlignment,
						bool								   DryRun,
						const ReclaimCallback&				   ChangeCallback	   = [](const MovedChunksArray&, const ChunkIndexArray&) {},
						const ClaimDiskReserveCallback&		   DiskReserveCallback = []() { return 0; });

	bool IterateChunks(const std::span<const BlockStoreLocation>& ChunkLocations, const IterateChunksCallback& Callback);

	bool IterateBlock(std::span<const BlockStoreLocation>	ChunkLocations,
					  std::span<const size_t>				ChunkIndexes,
					  const IterateChunksSmallSizeCallback& SmallSizeCallback,
					  const IterateChunksLargeSizeCallback& LargeSizeCallback,
					  uint64_t								LargeSizeLimit = 0);

	void CompactBlocks(
		const BlockStoreCompactState&	CompactState,
		uint32_t						PayloadAlignment,
		const CompactCallback&			ChangeCallback		= [](const MovedChunksArray&, uint64_t) { return true; },
		const ClaimDiskReserveCallback& DiskReserveCallback = []() { return 0; },
		std::string_view				LogPrefix			= {});

	inline uint64_t TotalSize() const { return m_TotalSize.load(std::memory_order::relaxed); }

private:
	static const char*			 GetBlockFileExtension();
	static std::filesystem::path GetBlockPath(const std::filesystem::path& BlocksBasePath, const uint32_t BlockIndex);
	uint32_t GetFreeBlockIndex(uint32_t StartProbeIndex, RwLock::ExclusiveLockScope&, std::filesystem::path& OutBlockPath) const;

	std::unordered_map<uint32_t, Ref<BlockStoreFile>> m_ChunkBlocks;

	mutable RwLock		  m_InsertLock;	 // used to serialize inserts
	Ref<BlockStoreFile>	  m_WriteBlock;
	std::uint32_t		  m_CurrentInsertOffset = 0;
	std::atomic_uint32_t  m_WriteBlockIndex{};
	std::vector<uint32_t> m_ActiveWriteBlocks;

	uint64_t			  m_MaxBlockSize  = 1u << 28;
	uint64_t			  m_MaxBlockCount = BlockStoreDiskLocation::MaxBlockIndex + 1;
	std::filesystem::path m_BlocksBasePath;

	std::atomic_uint64_t m_TotalSize{};
};

class BlockStoreCompactState
{
public:
	BlockStoreCompactState() = default;

	void IncludeBlocks(const BlockStore::BlockEntryCountMap& BlockEntryCountMap)
	{
		size_t EntryCountTotal = 0;
		for (auto& BlockUsageIt : BlockEntryCountMap)
		{
			uint32_t BlockIndex = BlockUsageIt.first;
			ZEN_ASSERT(m_BlockIndexToChunkMapIndex.find(BlockIndex) == m_BlockIndexToChunkMapIndex.end());

			m_KeepChunks.emplace_back(std::vector<size_t>());
			m_KeepChunks.back().reserve(BlockUsageIt.second);
			m_BlockIndexToChunkMapIndex.insert_or_assign(BlockIndex, m_KeepChunks.size() - 1);
			EntryCountTotal += BlockUsageIt.second;
		}
		m_ChunkLocations.reserve(EntryCountTotal);
	}

	void IncludeBlock(uint32_t BlockIndex)
	{
		if (m_BlockIndexToChunkMapIndex.find(BlockIndex) == m_BlockIndexToChunkMapIndex.end())
		{
			m_KeepChunks.emplace_back(std::vector<size_t>());
			m_BlockIndexToChunkMapIndex.insert_or_assign(BlockIndex, m_KeepChunks.size() - 1);
		}
	}

	bool AddKeepLocation(const BlockStoreLocation& Location)
	{
		auto It = m_BlockIndexToChunkMapIndex.find(Location.BlockIndex);
		if (It == m_BlockIndexToChunkMapIndex.end())
		{
			return false;
		}

		std::vector<size_t>& KeepChunks = m_KeepChunks[It->second];
		size_t				 Index		= m_ChunkLocations.size();
		KeepChunks.push_back(Index);
		m_ChunkLocations.push_back(Location);
		return true;
	};

	const BlockStoreLocation& GetLocation(size_t Index) const { return m_ChunkLocations[Index]; }

	void IterateBlocks(std::function<bool(uint32_t								 BlockIndex,
										  const std::vector<size_t>&			 KeepChunkIndexes,
										  const std::vector<BlockStoreLocation>& ChunkLocations)> Callback) const
	{
		for (auto It : m_BlockIndexToChunkMapIndex)
		{
			size_t ChunkMapIndex = It.second;
			bool   Continue		 = Callback(It.first, m_KeepChunks[ChunkMapIndex], m_ChunkLocations);
			if (!Continue)
			{
				break;
			}
		}
	}

private:
	std::unordered_map<uint32_t, size_t> m_BlockIndexToChunkMapIndex;  // Maps to which vector in BlockKeepChunks to use for a block
	std::vector<std::vector<size_t>>	 m_KeepChunks;				   // One vector per block index with index into ChunkLocations
	std::vector<BlockStoreLocation>		 m_ChunkLocations;
};

void blockstore_forcelink();

}  // namespace zen
